package com.boge.ai.utils;

import com.boge.ai.entity.embedd.EmbeddingData;
import com.boge.ai.entity.embedd.EmbeddingResponse;
import okhttp3.*;
import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.RealVector;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

public class TextSimilarityUtils {

    public static void main(String[] args) throws Exception {
        String query = "人工智能伦理";
        List<String> documents = new ArrayList<>();
        documents.add("全球科技巨头联合发布AI伦理指南，强调透明度与公平性");
        documents.add("欧盟通过新的法规，要求所有AI系统必须符合严格的隐私保护标准");
        documents.add("科学家警告：如果不加以控制，AI可能会加剧社会不平等");
        documents.add("国际会议讨论如何防止AI武器化，并呼吁制定国际条约");
        documents.add("某国政府宣布将投资数十亿美元用于支持可持续发展的AI技术研究");

        List<RealVector> docVecs = getEmbeddings(documents);
        RealVector queryVec = getEmbeddings(List.of(query)).get(0);

        System.out.println("Cosine distance:");
        System.out.println(cosSim(queryVec, queryVec));
        for (RealVector vec : docVecs) {
            System.out.println(cosSim(queryVec, vec));
        }

        System.out.println("\nEuclidean distance:");
        System.out.println(l2(queryVec, queryVec));
        for (RealVector vec : docVecs) {
            System.out.println(l2(queryVec, vec));
        }
    }

    public static List<RealVector> getEmbeddings(List<String> texts) throws Exception {
        List<RealVector> embeddings = new ArrayList<>();
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(20, TimeUnit.SECONDS)
                .readTimeout(20, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");


        for (String text : texts) {
            String requestBodyJson = String.format("""
                {
                  "input": "%s",
                  "model": "%s"
                }
                """, text, "text-embedding-3-small");

            RequestBody body = RequestBody.create(mediaType, requestBodyJson);
            Request request = new Request.Builder()
                    .url("https://api.openai-hk.com/v1/embeddings")
                    .method("POST", body)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("Accept", "application/json")
                    .addHeader("Authorization", "Bearer hk-w3q3id1000054953c829f33c861457133e3048fa220f1076")
                    .build();
            Response response = client.newCall(request).execute();
            ResponseBody responseBody = response.body();
            String jsonString = responseBody.string();
            EmbeddingResponse embeddingResponse = JsonToHashMapUtils.parseJsonToEmbeddingResponse(jsonString);
            if (embeddingResponse != null) {
                List<EmbeddingData> data = embeddingResponse.getData();
                if (data != null) {
                    for (EmbeddingData datum : data) {
                        double[] vector = datum.getEmbedding().stream().mapToDouble(Double::doubleValue).toArray();
                        embeddings.add(new ArrayRealVector(vector));
                    }
                }
            }
        }

        return embeddings;
    }



    public static List<List<Float>> getEmbeddingsFloat(List<String> texts) throws Exception {
        List<List<Float>> embeddings = new ArrayList<>();
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(20, TimeUnit.SECONDS)
                .readTimeout(20, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");

        for (String text : texts) {
            String requestBodyJson = String.format("""
                {
                  "input": "%s",
                  "model": "%s"
                }
                """, text, "text-embedding-3-small");

            RequestBody body = RequestBody.create(mediaType, requestBodyJson);
            Request request = new Request.Builder()
                    .url("https://api.openai-hk.com/v1/embeddings")
                    .method("POST", body)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("Accept", "application/json")
                    .addHeader("Authorization", "Bearer hk-w3q3id1000054953c829f33c861457133e3048fa220f1076")
                    .build();

            Response response = client.newCall(request).execute();
            if (!response.isSuccessful()) throw new RuntimeException("Unexpected code " + response);

            ResponseBody responseBody = response.body();
            if (responseBody == null) throw new RuntimeException("Empty response body");

            String jsonString = responseBody.string();
            EmbeddingResponse embeddingResponse = JsonToHashMapUtils.parseJsonToEmbeddingResponse(jsonString);

            if (embeddingResponse != null) {
                List<EmbeddingData> data = embeddingResponse.getData();
                if (data != null && !data.isEmpty()) {
                    // 获取第一个 embedding 向量（通常每个 input 返回一个）
                    List<Float> vector = new ArrayList<>();
                    for (Double d : data.get(0).getEmbedding()) {
                        vector.add(d.floatValue());
                    }
                    embeddings.add(vector);
                }
            }
        }

        return embeddings;
    }

    /**
     * 余弦距离 -- 越⼤越相似
     * @param a
     * @param b
     * @return
     */
    private static double cosSim(RealVector a, RealVector b) {
        return a.dotProduct(b) / (a.getNorm() * b.getNorm());
    }

    /**
     * 欧式距离 -- 越⼩越相似
     * @param a
     * @param b
     * @return
     */
    private static double l2(RealVector a, RealVector b) {
        return a.subtract(b).getNorm();
    }
}
